package com.fanshuai.algorithms.greedy;

import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;

import java.util.*;

/**
 * 霍夫曼编码使用变长编码表对源符号（如文件中的一个字母）进行编码，其中变长编码表是通过一种评估来源符号出现机率的方法得到的，
 * 出现机率高的字母使用较短的编码，反之出现机率低的则使用较长的编码，这便使编码之后的字符串的平均长度、期望值降低，从而达到无损压缩数据的目的。
 */
public class Huffman {
    static char zero = '\0';
    /**
     * Huffman树节点
     */
    static class HuffmanNode {
        public char ch; //字符
        public int count;  //次数
        public HuffmanNode left, right;

        public HuffmanNode(char ch, int count) {
            this.ch = ch;
            this.count = count;
        }

        @Override
        public String toString() {
            return "HuffmanNode{" +
                    "ch=" + ch +
                    ", count=" + count +
                    '}';
        }
    }

    /**
     * 根据字符及其出现的频率，构造huffman树
     * @param nodes
     * @return
     */
    private static HuffmanNode buildHuffmanTree(List<HuffmanNode> nodes) {
        if (CollectionUtils.isEmpty(nodes)) {
            return null;
        }

        HuffmanNode root = null;
        while (nodes.size() > 1) {
            HuffmanNode[] min = find2Min(nodes);
            nodes.remove(min[0]);
            nodes.remove(min[1]);

            HuffmanNode left = min[0];
            HuffmanNode right = min[1];
            HuffmanNode parent = new HuffmanNode(zero, min[0].count + min[1].count);
            parent.left = left;
            parent.right = right;
            nodes.add(parent);

            if (nodes.size() == 1) {
                root = parent;
            }
        }

        return root;
    }

    /**
     * 获取出现次数最小的两个字符
     * @param chars
     * @return
     */
    private static HuffmanNode[] find2Min(List<HuffmanNode> chars) {
        HuffmanNode min1 = null;
        HuffmanNode min2 = null;

        int minCount = Integer.MAX_VALUE;
        for (HuffmanNode node : chars) {
            if (node.count < minCount) {
                minCount = node.count;
                min1 = node;
            }
        }

        int min2Count = Integer.MAX_VALUE;
        for (HuffmanNode node : chars) {
            if (node.count < min2Count && node != min1) {
                min2Count = node.count;
                min2 = node;
            }
        }

        return new HuffmanNode[] {min1, min2};
    }

    /**
     * 根据huffman树生成huffman编码
     *
     * 使用回溯法遍历，遇到有效字母的叶节点，则记录该字母的编码为从根节点到当前节点的路径
     * @param huffmanNode
     * @return
     */
    private static Map<Character, String> generateHuffmanCode(HuffmanNode huffmanNode) {
        Map<Character, String> codeMap = new HashMap<>();
        generateHuffmanCode(huffmanNode, new StringBuilder(), codeMap);

        return codeMap;
    }

    private static void generateHuffmanCode(HuffmanNode node, StringBuilder builder, Map<Character, String> codeMap) {
        if (null == node) {
            return;
        }
        if (node.left == null && node.right == null) {//递归出口
            if (node.ch != zero) { //记录有效节点
                String code = builder.toString();
                codeMap.put(node.ch, code);
            }
            return;
        }

        builder.append("0");
        generateHuffmanCode(node.left, builder, codeMap);
        builder.deleteCharAt(builder.length() - 1);

        builder.append("1");
        generateHuffmanCode(node.right, builder, codeMap);
        builder.deleteCharAt(builder.length() - 1);
    }

    public static Map<Character, String> huffmanEncode(String s) {
        if (StringUtils.isBlank(s)) {
            return new HashMap<>();
        }

        List<HuffmanNode> list = new ArrayList<>();

        //统计字符出现次数
        Map<Character, Integer> countMap = new HashMap<>();
        for (char ch : s.toCharArray()) {
            if (Character.isWhitespace(ch)) {
                continue;
            }

            if (countMap.containsKey(ch)) {
                countMap.put(ch, countMap.get(ch) + 1);
            } else {
                countMap.put(ch, 1);
            }
        }
        for (char ch : countMap.keySet()) {
            HuffmanNode node = new HuffmanNode(ch, countMap.get(ch));
            list.add(node);
        }

        //构造huffman树
        HuffmanNode root = buildHuffmanTree(list);
        //根据huffman树生成编码
        return generateHuffmanCode(root);
    }

    public static void main(String[] args) {
        String s = "BABACAC ADADABB CBABEBE DDABEEEBB";

        Map<Character, String> huffmanCode = huffmanEncode(s);
        System.out.println(huffmanCode);
    }
}
